import os
import json

class CLUENERWrapper:
    def __init__(self, data_root="data/cluener/", data_type="train") -> None:
        
        self.id2type = self.load_id2type(os.path.join(data_root, "id2type.txt"))
        self.type2id = {k:v for v,k in self.id2type.items()}
        
        self._data_type = data_type
        self._data_filename = os.path.join(data_root, self._data_type+".jsonl")
        
        with open(self._data_filename, "r", encoding="utf-8") as fp:
            self._data = fp.readlines()
        
        pass

    def __len__(self):
        return len(self._data)
    
    def __getitem__(self, index):
        jsonl = self._data[index]
        jsond = json.loads(jsonl)
        text = jsond["text"]
        label = []
        
        for ent_type, ent_info in jsond["label"].items():
            type_id = self.type2id[ent_type]  # 将文本转为索引

            for ent_str, span_list in ent_info.items():  # 同样的字符串实体可能在1个文本中出现多次, 序列标注不需要用到实体的文本形式，只用到span索引
                label.extend([(type_id, (span[0], span[1]))for span in span_list])

        return {"text":text, "label":label}
    
    @staticmethod
    def load_id2type(filename="id2type.txt"):
        
        type_list = None
        with open(filename, "r", encoding="utf-8") as fp:
            type_list = fp.readlines()
        
        id2type = {k:v.strip("\n\r\t ") for k,v in enumerate(type_list)}
        
        return id2type